package prod_cons;

import java.util.Random;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicIntegerArray;

public class Pow2Container implements TaskContainer {
	int consumersNum;
	ConcurrentLinkedQueue<Runnable>[] queues;
	AtomicIntegerArray queueSizes;
	
	@SuppressWarnings("unchecked")
	public void setConsumersNum(int consumersNum) {
		this.consumersNum = consumersNum;
		queues = new ConcurrentLinkedQueue[consumersNum];
		queueSizes = new AtomicIntegerArray(consumersNum);
		for(int i = 0; i < consumersNum; i++) {
			queues[i] = new ConcurrentLinkedQueue<Runnable>();
			queueSizes.set(i, 0);
		}
	}
	
	private static final ThreadLocal<Integer> queueNum = new ThreadLocal<Integer>() {
		protected Integer initialValue() {
			return new Integer(0);
		}
	};
	
	AtomicInteger consumerCount = new AtomicInteger(0);
	public void registerConsumer() {
		queueNum.set(consumerCount.getAndIncrement());
	}
	
	Random random = new Random();
	
	@Override
	public void put(Runnable task) {
		//take two random choices, chose the queue with the least number of tasks
		int q1Idx = random.nextInt(this.consumersNum);
		int q2Idx = random.nextInt(this.consumersNum);
		ConcurrentLinkedQueue<Runnable> q1 = queues[q1Idx];
		ConcurrentLinkedQueue<Runnable> q2 = queues[q2Idx];
		if (queueSizes.get(q1Idx) < queueSizes.get(q2Idx)) {
			q1.offer(task);
			queueSizes.incrementAndGet(q1Idx);
		} else {
			q2.offer(task);
			queueSizes.incrementAndGet(q2Idx);
		}
	}

	@Override
	public Runnable get() {
		int myQueue = queueNum.get();
		Runnable task = queues[myQueue].poll();
		if (task != null) {
			queueSizes.decrementAndGet(myQueue);
			return task;
		}
		
		for(int i = 1; i < this.consumersNum; i++) {
			int queueIdx = (myQueue+i)%consumersNum; 
			task = queues[queueIdx].poll();
			if (task != null) {
				queueSizes.decrementAndGet(queueIdx);
				return task;
			}
		}
		return null;
	}
}